import os

import dotenv
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.memory import ChatMessageHistory
from langchain_community.document_loaders import CSVLoader
from langchain_community.tools import TavilySearchResults
from langchain_community.vectorstores import Chroma
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory, RunnableConfig
from langchain_core.tools import create_retriever_tool, Tool
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langsmith import Client

# 创建历史记录保存会话
session = dict()

# 创建和获取上下文记忆
def get_history_session(session_id) -> BaseChatMessageHistory:
    if session_id not in session:
        session[session_id] = ChatMessageHistory()
    return session[session_id]


if __name__ == "__main__":

    # 加载文档
    loader = CSVLoader("../assest/LangChain.md", encoding="utf-8")
    docs = loader.load()

    # 拆分文档
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    chunks = text_splitter.split_documents(docs)

    # 创建嵌入模型
    dotenv.load_dotenv()

    os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
    os.environ['OPENAI_BASE_URL'] = os.getenv("OPENAI_BASE_URL")

    embed_model = OpenAIEmbeddings(model="text-embedding-3-large")

    # 添加到向量库
    chroma = Chroma.from_documents(chunks, embed_model)

    # 创建天气检索工具
    os.environ['TAVILY_API_KEY'] = os.getenv("TAVILY_API_KEY")
    search = TavilySearchResults(max_results=20)

    search_tool = Tool(
        name="search_tool",
        description="查询天气资讯",
        func=search.run,
    )

    # 创建知识库召回工具
    retriever = chroma.as_retriever(
        search_type="similarity_score_threshold",
        search_kwargs={"k": 10, "score_threshold": 0.5}
    )

    retriever_tool = create_retriever_tool(retriever, "knowledge_search", "知识库搜索")

    # 创建大模型实例
    llm = ChatOpenAI(model="gpt-4o-mini")

    # 创建工具集
    tools = [search_tool, retriever_tool]

    # 获取在线提示词模板
    client = Client(api_key=os.getenv("PROMPT_API_KEY"))
    prompt = client.pull_prompt("hwchase17/openai-tools-agent", include_model=True)

    # 创建agent
    agent = create_tool_calling_agent(llm, tools, prompt)

    # 创建执行器
    executor = AgentExecutor(tools=tools, agent=agent, verbose=True)

    # 上下文记忆添加到agent执行器
    executor_with_history = RunnableWithMessageHistory(
        runnable=executor,
        get_session_history=get_history_session,
        input_messages_key="input",
        history_messages_key="chat_history"
    )

    # 调用执行器
    executor_with_history.invoke(
        {"input": "杭州明天的天气如何？"},
        config=RunnableConfig(configurable={"session_id": "123"})
    )

    # 测试历史对话记录
    results = executor_with_history.invoke(
        {"input": "上海呢？"},
        config=RunnableConfig(configurable={"session_id": "123"})
    )

    print(results)





